import cv2
import numpy as np

# 角点检测参数
feature_params = dict(maxCorners = 100,
                      qualityLevel = 0.3,
                      minDistance = 7)
# lucas kanade光流法参数
lk_params = dict(winSize = (15, 15),
                 maxLevel = 2)
cap = cv2.VideoCapture(r'D:\opencv\sources\samples\data\vtest.avi')

# 拿到第一帧图像
ret, prev = cap.read()
prevGray = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY)
# 获取图像中最好的角点特征,None表示在整幅图上寻找角点。
p0 = cv2.goodFeaturesToTrack(prevGray, mask=None, **feature_params)

while (True):
    ret, frame = cap.read()
    if not ret:
        break
    frameGray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    # 计算光流
    p1, st, err = cv2.calcOpticalFlowPyrLK(prevGray, frameGray, p0, None, **lk_params)
    # 选取好的跟踪点
    goodPoints = p1[st == 1]
    goodPrevPoints = p0[st == 1]
    # 在结果图像中迭加画出特征点和计算出来的光流向量
    res = frame.copy()
    drawColor = (0, 0, 255)
    for i, (cur, prev) in enumerate(zip(goodPoints, goodPrevPoints)):
        x0, y0 = cur.ravel()
        x1, y1 = prev.ravel()
        cv2.line(res, (x0, y0), (x1, y1), drawColor)
        cv2.circle(res, (x0, y0), 5, drawColor)

    #更新上一帧
    prevGray = frameGray.copy()
    p0 = goodPoints.reshape(-1, 1, 2)

    #显示计算结果图像
    cv2.imshow('frame', res)

    k = cv2.waitKey(30)
    if k == 27:
        break


cv2.destroyAllWindows()
cap.release()